import os
import sys

import numpy as np
import pandas as pd

import cv2
from PIL import Image
import torch
from tqdm import tqdm
import ttach as tta

# refrence
#  > https://www.kaggle.com/raufyagfarov/mmdetection-yolov4-pipeline-with-tta

modelPath = '    '

model = torch.load(modelPath)
output = []

transforms = tta.Compose(
    [
        tta.HorizontalFlip(),
        tta.Rotate90(angles=[0, 90, 180, 270]),
    ]
)



def ttaInfer(model, imageSets):

    # transformer : flips + rotation 0, 90, 180, 270
    # merge_mod : mean, gmean (geometric mean), sum, max, min, tsharpen (temperature sharpen with t=0.5)
    ttaModel = tta.SegmentationTTAWrapper(
        model, tta.aliases.d4_transform(), merge_mode='tsharpen')

    for image in imageSets:
        outputs = ttaModel(imageSets)

